from matplotlib import pyplot
import minuit
import numpy

import histogram

import ypipi


def gauss(x, A, mu, sigma):
    return A*numpy.exp(-0.5*((x - mu)/sigma)**2)

def bg(x, p0, p1, p2, p3, p4, p5, p6):
    return numpy.polynomial.chebyshev.chebval(x, [p0, p1, p2, p3, p4, p5, p6])

def y_calc(x, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6):
    return gauss(x, A, mu, sigma) + bg(x, p0, p1, p2, p3, p4, p5, p6)

def chisq(A, mu, sigma, p0, p1, p2, p3, p4, p5, p6):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (y_calc(_g_x, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print _g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)

def bg_chisq(p0, p1, p2, p3, p4, p5, p6):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (bg(_g_x, p0, p1, p2, p3, p4, p5, p6) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print 'bg_chisq:',_g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)

def sig_chisq(A, mu, sigma):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (gauss(_g_x, A, mu, sigma) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print 'sig_chisq:',_g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)

def m_mumu_fit(path, fit_range=(9.35, 9.7)):
    hsig = histogram.load(path)
    h = hsig['m_mumu']

    lower_limit = h.bin_index(fit_range[0])
    upper_limit = h.bin_index(fit_range[1])

    global _g_x, _g_y, _g_yunc, _g_ncalls
    
    _g_y = h.bins()[lower_limit:upper_limit]

    x_full = numpy.array(map(h.bin_center, range(h.nbins)))
    _g_x = x_full[lower_limit:upper_limit]

    _g_yunc = numpy.sqrt(_g_y)
    _g_yunc = numpy.where(_g_yunc > 0, _g_yunc, numpy.ones(len(_g_yunc)))

    _g_ncalls = 0
    
    m = minuit.Minuit(sig_chisq)

    m.values['A'] = 10
    m.values['mu'] = 9.5
    m.values['sigma'] = 0.01

    initial_values = {}
    initial_values.update(m.values)

    try:
        m.migrad()
    except minuit.MinuitError:
        print "Died on signal fit"
        raise
    
    figx = 8
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.8, 0.8))

    info = [m.values['A'], m.errors['A'],
            m.values['mu'], m.errors['mu'],
            m.values['sigma'], m.errors['sigma'],
            sig_chisq(**initial_values) / (len(_g_x) - len(m.values.keys())), 
            sig_chisq(**m.values) / (len(_g_x) - len(m.values.keys())), 
            ]

    fptex = ypipi.textable.write_sig_table(info)

    pyplot.text(0.95, 0.95,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF', 'alpha': '0.7'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = gauss(xfit, **m.values)

    pyplot.plot(xfit, yfit, 'b-')
    pyplot.xlabel(r"M($\mu^+\mu^-$) GeV")
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title)

    path_out = h.name + '_m_mumufit.png'
    pyplot.savefig(path_out)

    path_out = h.name + '_m_mumufit.pdf'
    pyplot.savefig(path_out)

    return m


def fixed_sig_fit(path, fit_range=(9.35, 9.7)):
    hdict = histogram.load(path)
    h = hdict['mm_off_pi0pi0']

    global _g_x, _g_y, _g_yunc, _g_ncalls

    lower_limit = h.bin_index(fit_range[0])
    upper_limit = h.bin_index(fit_range[1])
    
    _g_y = h.bins()[lower_limit:upper_limit]
    x_full = numpy.array(map(h.bin_center, range(h.nbins)))
    _g_x = x_full[lower_limit:upper_limit]

    _g_yunc = numpy.sqrt(numpy.fabs(_g_y))
    _g_yunc = numpy.where(_g_yunc > 0, _g_yunc, numpy.ones(len(_g_yunc)))

    m = minuit.Minuit(chisq)

    m.values['A'] = 3.04
    m.values['mu'] = 9.5
    m.values['sigma'] = 0.0313

    m.values['p0'] = -33.4
    m.values['p1'] = -3.54
    m.values['p2'] = 0.712
    m.values['p3'] = -0.177
    
    m.fixed['mu'] = True
    m.fixed['sigma'] = True

    initial_values = {}
    initial_values.update(m.values)

    _g_ncalls = 0
    
    try:
        m.migrad()
    except minuit.MinuitError:
        pass
    
    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.13, 0.6, 0.75))

    info = [m.values['A'], m.errors['A'],
            m.values['mu'], m.errors['mu'],
            m.values['sigma'], m.errors['sigma'],
            m.values['p0'], m.errors['p0'],
            m.values['p1'], m.errors['p1'],
            m.values['p2'], m.errors['p2'],
            m.values['p3'], m.errors['p3'],
            ]

    initial_chi2 = chisq(**initial_values) / (len(_g_x) - len(m.values.keys()))
    final_chi2 = chisq(**m.values) / (len(_g_x) - len(m.values.keys()))

    fptex = ypipi.textable.write_minuit_table(m, initial_chi2=initial_chi2, final_chi2=final_chi2,
                                              param_names='A mu sigma p0 p1 p2 p3'.split(),
                                              param_latex={'A': 'A', 'mu':'\mu', 'sigma':'\sigma',
                                               'p0':'p_0', 'p1': 'p_1', 'p2': 'p_2',
                                               'p3':'p_3', 'p4': 'p_4', 'p5': 'p_5',}
                                               )

    pyplot.text(1.35, 0.75,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = y_calc(xfit, **m.values)
    pyplot.plot(xfit, yfit, 'b-')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title)

    path_out = h.name + '_sigfixed_fit.png'
    pyplot.savefig(path_out)
    path_out = h.name + '_sigfixed_fit.pdf'
    pyplot.savefig(path_out)

    return m, initial_values, initial_chi2, final_chi2


def plot_fit_bg_subtracted(path, m, initial_values, chi2i, chi2f, fit_range=(9.35, 9.7)):
    h = histogram.load(path)['mm_off_pi0pi0']

    lower_limit = h.bin_index(fit_range[0])
    upper_limit = h.bin_index(fit_range[1])
        
    x = numpy.array([h.bin_center(i) for i in xrange(h.nbins)])
    y = h.bins() - bg(x, 
                      m.values['p0'], 
                      m.values['p1'], 
                      m.values['p2'], 
                      m.values['p3'], 
                      m.values['p4'], 
                      m.values['p5'], 
                      m.values['p6'], 
                      )

    x = x[lower_limit:upper_limit]
    y = y[lower_limit:upper_limit]
    
    yunc = numpy.where(y > 0, numpy.sqrt(y), numpy.ones(len(y)))

    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.15, 0.141, 0.6, 0.75))

    fptex = ypipi.textable.write_minuit_table(m, chi2i, chi2f, 
                                        param_names='A mu sigma'.split(),
                                        param_latex={'A':  'A', 
                                                     'mu': '\mu',
                                                     'sigma': '\sigma'}
                                        )

    pyplot.text(1.35, 0.95,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(x, y, yerr=yunc, fmt='r.')

    xfit = numpy.linspace(min(x), max(x),
                          10000)

    yfit = gauss(xfit, m.values['A'], m.values['mu'], m.values['sigma'])
    pyplot.plot(xfit, yfit, 'b-')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title)

    path_out = h.name + '_subtracted.png'
    pyplot.savefig(path_out)
    path_out = h.name + '_subtracted.pdf'
    pyplot.savefig(path_out)


def bg_subtracted_fit(m_with_bg, path, fit_range=(9.35, 9.7)):
    hdict = histogram.load(path)
    h = hdict['mm_off_pi0pi0']
    
    global _g_x, _g_y, _g_yunc, _g_ncalls

    lower_limit = h.bin_index(fit_range[0])
    upper_limit = h.bin_index(fit_range[1])
    
    x_full = numpy.array(map(h.bin_center, range(h.nbins)))
    _g_x = x_full[lower_limit:upper_limit]

    bgvals = bg(_g_x, 
                m_with_bg.values['p0'], 
                m_with_bg.values['p1'], 
                m_with_bg.values['p2'], 
                m_with_bg.values['p3'], 
                m_with_bg.values['p4'], 
                m_with_bg.values['p5'], 
                m_with_bg.values['p6'], 
                )

    _g_y = h.bins()[lower_limit:upper_limit]

    _g_y -= bgvals

    _g_yunc = numpy.sqrt(numpy.fabs(_g_y))
    _g_yunc = numpy.where(_g_yunc > 0, _g_yunc, numpy.ones(len(_g_yunc)))
    
    pyplot.clf()
    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')
    pyplot.savefig("subtracted_signal.png")

    m = minuit.Minuit(sig_chisq)

    m.values['A']     = m_with_bg.values['A']
    m.values['mu']    = m_with_bg.values['mu']
    m.values['sigma'] = m_with_bg.values['sigma']

    initial_values = {}
    initial_values.update(m.values)

    m.initial_values = initial_values

    try:
        m.migrad()
    except minuit.MinuitError:
        pass
    
    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.13, 0.6, 0.75))

    initial_chi2 = sig_chisq(**initial_values) / (len(_g_x) - len(m.values.keys()))
    final_chi2   = sig_chisq(**m.values) / (len(_g_x) - len(m.values.keys())) 

    fptex = ypipi.textable.write_minuit_table(m, initial_chi2, final_chi2, 
                                              param_names='A mu sigma'.split(),
                                              param_latex={'A':  'A', 
                                                           'mu': '\mu',
                                                           'sigma': '\sigma'}
                                                           )

    pyplot.text(1.35, 0.75,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = gauss(xfit, **m.values)
    pyplot.plot(xfit, yfit, 'b-')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title)

    path_out = h.name + '_subtracted_fit.png'
    pyplot.savefig(path_out)
    path_out = h.name + '_subtracted_fit.pdf'
    pyplot.savefig(path_out)
    
